{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "from langchain.schema import Document\n",
    "import fitz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "class tokenizer():\n",
    "    def __init__(self) -> None:\n",
    "        pass\n",
    "    @classmethod\n",
    "    def tokenize(self, text):\n",
    "        return list(text)\n",
    "pdf_path = '/home/featurize/work/llm-rag/assets/初赛训练数据集.pdf'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def split_by_blocks(block_size: int = 500, block_overlap: int = 125) -> List[Document]:\n",
    "        \"\"\"根据\n",
    "\n",
    "        Args:\n",
    "            block_size (int, optional): _description_. Defaults to 500.\n",
    "            block_overlap (int, optional): _description_. Defaults to 125.\n",
    "\n",
    "        Returns:\n",
    "            List[Document]: _description_\n",
    "        \"\"\"\n",
    "        block_texts = []\n",
    "        with fitz.open(pdf_path) as pdf:\n",
    "            for page in pdf:\n",
    "                page_block_texts = []\n",
    "                blocks = page.get_text('blocks')\n",
    "                for i, block in enumerate(blocks):\n",
    "                    block_text = block[4][:-1] # 去掉\\n\n",
    "                    if block_text.isdigit():\n",
    "                        # 去掉页码\n",
    "                        continue\n",
    "                    if '----' in block_text or '....' in block_text:\n",
    "                        # 去掉目录\n",
    "                        continue\n",
    "                    page_block_texts.append(block_text)\n",
    "                tables = list(page.find_tables())\n",
    "                if tables:\n",
    "                    # 找到table对应的block替换为模型可读顺序文本\n",
    "                    for table in tables:\n",
    "                        header = table.extract()[0]\n",
    "                        header = [str(h) for h in header]\n",
    "                        header_text = '\\n'.join(header)\n",
    "                        num_table_rows = table.row_count\n",
    "                        table_text = '以下是一张表格:表格的表头为:{}'.format(header)\n",
    "                        table_text += '表格的信息如下:\\n'\n",
    "                        for table_row in table.extract()[1:]:\n",
    "                            row_text = ''\n",
    "                            for row_item in zip(header, table_row):\n",
    "                                row_text += '{}:{}'.format(row_item[0], row_item[1])\n",
    "                            table_text += row_text\n",
    "                        table_text += '\\n'\n",
    "                        for i, page_block_text in enumerate(page_block_texts):\n",
    "                            if header_text in page_block_text:\n",
    "                                table_start_idx = i\n",
    "                                table_end_idx = i + num_table_rows\n",
    "                                break\n",
    "                        page_block_texts[table_start_idx: table_end_idx] = [table_text]\n",
    "                block_texts.extend(page_block_texts)\n",
    "        docs = []\n",
    "        document_text = ''\n",
    "        for block_text in block_texts:\n",
    "            num_previous_tokens = len(tokenizer.tokenize(document_text))\n",
    "            num_block_tokens = len(tokenizer.tokenize(block_text))\n",
    "            if num_block_tokens > block_size:\n",
    "                docs.append(Document(page_content=document_text, metadata={'num_tokens': num_previous_tokens}))\n",
    "                document_text = ''\n",
    "                docs.append(Document(page_content=block_text, metadata={'num_tokens': num_block_tokens}))\n",
    "                continue\n",
    "            if (num_previous_tokens + num_block_tokens) >= block_size:\n",
    "                doc = Document(page_content=document_text, metadata={'num_tokens': len(tokenizer.tokenize(document_text))})\n",
    "                docs.append(doc)\n",
    "                document_text = ''\n",
    "                continue\n",
    "            document_text += block_text\n",
    "        return docs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "docs=split_by_blocks()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Document(page_content='欢迎\\n感谢您选择了具有优良安全性、舒适性、动力性和经济性的Lynk & Co领克汽车。首次使用前请仔细、完整地阅读本手册内容，将有助于您更好地了解和使用车辆。本手册中的所有资料均为出版时的最新资料，但本公司将对产品进行不断的改进和优化，您所购的车辆可能与本手册中的描述有所不同，请以实际\\n接收的车辆为准。如您有任何问题，或需要预约服务，请拨打电话4006-010101 联系我们。您也可以开车前往Lynk & Co领克中心。在抵达之前，请您注意驾车安全。© 领克汽车销售有限公司目录前言用车前准备装载货物上车和下车驾驶前的准备仪表和灯光目录安全出行启动和驾驶目录驾驶辅助泊车空调中央显示屏目录Lynk & Co App高压系统保养和维护OTA升级紧急情况下技术资料目录前言前言本手册相关的重要信息\\n领克汽车销售有限公司（下称“Lynk & Co领克”）建议您在首次使用\\n车辆前，认真阅读本手册内容。为了更好地理解本手册中的内容，您\\n需要了解以下所有信息。提示信息<image: DeviceRGB, width: 54, height: 54, bpc: 8警告！人身伤害', metadata={'num_tokens': 486})"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "docs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
